import mindspore
import numpy as np
from mindspore import context
from src.nnunet.training.network_training.nnUNetTrainerV2 import nnUNetTrainerV2
context.set_context(mode=context.GRAPH_MODE, device_target="GPU", save_graphs=False, enable_graph_kernel=False)
from src.nnunet.training.model_restore import load_model_and_checkpoint_files
model = "/home/wanyi/work_space_lzy/nnunet-master/src/nnUNetFrame/DATASET/nnUNet_trained_models/nnUNet/2d/Task004_Hippocampus/nnUNetTrainerV2__nnUNetPlansv2.1"
checkpoint_name = 'model_final_checkpoint'

trainer, params = load_model_and_checkpoint_files(model, folds=0, mixed_precision=False,
                                                      checkpoint_name=checkpoint_name)
#plans_file =
#fold =
init=('/home/wanyi/work_space_lzy/nnunet-master/src/nnUNetFrame/DATASET/nnUNet_preprocessed/Task004_Hippocampus/nnUNetPlansv2.1_plans_2D.pkl', 0, '/home/wanyi/work_space_lzy/nnunet-master/src/nnUNetFrame/DATASET/nnUNet_trained_models/nnUNet/2d/Task004_Hippocampus/nnUNetTrainerV2__nnUNetPlansv2.1', '/home/wanyi/work_space_lzy/nnunet-master/src/nnUNetFrame/DATASET/nnUNet_preprocessed/Task004_Hippocampus', True, 0, True, True, True)

trainer2 = nnUNetTrainerV2(*init)
trainer2.initialize()
trainer2.initialize_network()
trainer2.network.set_train(True)


data = np.zeros(shape=(366, 1, 56, 40), dtype=np.float32)
data = mindspore.Tensor(data)
trainer.network.set_train(True)
#print(trainer2.network(data))
print(trainer.network(data))


